import os
from dotenv import load_dotenv
load_dotenv() # 加载.env文件中的环境变量
import numpy as np
from datasets import Dataset
from ragas.metrics import Faithfulness, AnswerRelevancy
from ragas.llms import LangchainLLMWrapper
from ragas.embeddings import LangchainEmbeddingsWrapper
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_huggingface import HuggingFaceEmbeddings
from sentence_transformers import SentenceTransformer
from ragas import evaluate
import logging

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)



# 准备评估用的LLM
try:
    llm = LangchainLLMWrapper(ChatOpenAI(
        model_name="qwen-plus", 
        openai_api_key="xxxxxxxx", 
        openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
    ))
    logger.info("Quen LLM初始化成功")
except Exception as e:
    logger.error(f"LLM初始化失败: {e}")
    raise

# 创建自定义的Embedding包装器来适配RAGAS
class LocalEmbeddings:    
    def __init__(self, embedding_function):
        self.embedding_function = embedding_function
    
    def embed_documents(self, texts):
        return self.embedding_function(texts)
    
    def embed_query(self, text):
        return self.embedding_function([text])[0]
    
    async def aembed_documents(self, texts):
        return self.embed_documents(texts)
    
    async def aembed_query(self, text):
        return self.embed_query(text)



# 准备数据集
# 这个数据集包含了问题、生成的答案以及相关的上下文信息
data = {
    "question": [
        "Who is the main character in Black Myth: Wukong?",
        "What are the special features of the combat system in Black Myth: Wukong?",
        "How is the visual quality of Black Myth: Wukong?",
    ],
    "answer": [
        "The main character in Black Myth: Wukong is Sun Wukong, based on the Chinese classic 'Journey to the West' but with a new interpretation. This version of Sun Wukong is more mature and brooding, showing a different personality from the traditional character.",
        "Black Myth: Wukong's combat system combines Chinese martial arts with Soulslike game features, including light and heavy attack combinations, technique transformations, and magic systems. Notably, Wukong can transform between different weapon forms during combat, such as his iconic staff and nunchucks, and use various mystical abilities.",
        "Black Myth: Wukong is developed using Unreal Engine 5, showcasing stunning visual quality. The game's scene modeling, lighting effects, and character details are all top-tier, particularly in its detailed recreation of traditional Chinese architecture and mythological settings.",
    ],
    "contexts": [
        [
            "Black Myth: Wukong is an action RPG developed by Game Science, featuring Sun Wukong as the protagonist based on 'Journey to the West' but with innovative interpretations. In the game, Wukong has a more composed personality and carries a special mission.",
            "The game is set in a mythological world, telling a new story that presents a different take on the traditional Sun Wukong character."
        ],
        [
            "The game's combat system is heavily influenced by Soulslike games while incorporating traditional Chinese martial arts elements. Players can utilize different weapon forms, including the iconic staff and other transforming weapons.",
            "During combat, players can unleash various mystical abilities, combined with light and heavy attacks and combo systems, creating a fluid and distinctive combat experience. The game also features a unique transformation system."
        ],
        [
            "Black Myth: Wukong demonstrates exceptional visual quality, built with Unreal Engine 5, achieving extremely high graphical fidelity. The game's environments and character models are meticulously crafted.",
            "The lighting effects, material rendering, and environmental details all reach AAA-level standards, perfectly capturing the atmosphere of an Eastern mythological world."
        ]
    ]
}

# 将字典转换为Hugging Face的Dataset对象，方便Ragas处理
dataset = Dataset.from_dict(data)

# 初始化本地embedding模型 - 使用项目标准化的local_embedding_function
# 验证本地embedding模型路径
bge_m3_model = "D:\\code\\aicode\\sentence-transformers\\bge-m3"
all_minilm_model = "D:\\code\\aicode\\sentence-transformers\\all-MiniLM-L6-v2"

# 检查模型路径
if not os.path.exists(bge_m3_model):
    logger.warning(f"BGE-M3模型路径不存在: {bge_m3_model}")
if not os.path.exists(all_minilm_model):
    logger.warning(f"all-MiniLM-L6-v2模型路径不存在: {all_minilm_model}")

# 初始化两个embedding模型
bge_m3_embedding_model = SentenceTransformer(bge_m3_model)
all_minilm_embedding_model = SentenceTransformer(all_minilm_model)

def create_local_embedding_function(model):
    """创建本地嵌入函数包装器 - 遵循项目规范"""
    def local_embedding_function(texts):
        if isinstance(texts, str):
            texts = [texts]
        embeddings = model.encode(texts, 
                                batch_size=12, 
                                convert_to_tensor=False,
                                normalize_embeddings=True)
        return embeddings.tolist()
    return local_embedding_function

# 创建两个embedding函数
bge_m3_embedding_function = create_local_embedding_function(bge_m3_embedding_model)
all_minilm_embedding_function = create_local_embedding_function(all_minilm_embedding_model)

print("\n=== Ragas评估指标说明 ===")
print("\nAnswerRelevancy（答案相关性）")
print("- 评估生成的答案与问题的相关程度")
print("- 使用embedding模型计算语义相似度")
print("- 将对比BGE-M3和all-MiniLM-L6-v2两个模型的相关性评分")

# 创建两个embedding包装器
bge_m3_embedding = LocalEmbeddings(bge_m3_embedding_function)
all_minilm_embedding = LocalEmbeddings(all_minilm_embedding_function)

# 创建答案相关性评估指标 - BGE-M3模型
print(f"\n=== 使用BGE-M3模型进行相关性评估 ===")
bge_m3_relevancy = [AnswerRelevancy(llm=llm, embeddings=bge_m3_embedding)]
bge_m3_result = evaluate(dataset, bge_m3_relevancy)
bge_m3_scores = bge_m3_result['answer_relevancy']
bge_m3_mean = np.mean(bge_m3_scores) if isinstance(bge_m3_scores, (list, np.ndarray)) else bge_m3_scores
print(f"BGE-M3 相关性评分: {bge_m3_mean:.4f}")
print(f"BGE-M3 各问题评分: {bge_m3_scores}")

# 创建答案相关性评估指标 - all-MiniLM-L6-v2模型
print(f"\n=== 使用all-MiniLM-L6-v2模型进行相关性评估 ===")
all_minilm_relevancy = [AnswerRelevancy(llm=llm, embeddings=all_minilm_embedding)]
all_minilm_result = evaluate(dataset, all_minilm_relevancy)
all_minilm_scores = all_minilm_result['answer_relevancy']
all_minilm_mean = np.mean(all_minilm_scores) if isinstance(all_minilm_scores, (list, np.ndarray)) else all_minilm_scores
print(f"all-MiniLM-L6-v2 相关性评分: {all_minilm_mean:.4f}")
print(f"all-MiniLM-L6-v2 各问题评分: {all_minilm_scores}")

# 对比两个模型的结果
print(f"\n=== 模型对比结果 ===")
print(f"BGE-M3 平均相关性评分: {bge_m3_mean:.4f}")
print(f"all-MiniLM-L6-v2 平均相关性评分: {all_minilm_mean:.4f}")
print(f"评分差异: {abs(bge_m3_mean - all_minilm_mean):.4f}")

# 详细对比每个问题的评分
print(f"\n=== 各问题详细对比 ===")
for i, question in enumerate(data["question"]):
    print(f"问题 {i+1}: {question[:50]}...")
    print(f"  BGE-M3 评分: {bge_m3_scores[i]:.4f}")
    print(f"  all-MiniLM-L6-v2 评分: {all_minilm_scores[i]:.4f}")
    print(f"  评分差异: {abs(bge_m3_scores[i] - all_minilm_scores[i]):.4f}")



